from tortoise import fields

from .admin import User
from .base import BaseModel, TimestampMixin
from app.schemas.models import AnotationStatusEnum, DataPartitionEnum, DataTypeEnum, IntelligenceModelTaskType, TaskStatusEnum


class IntelligenceModel(BaseModel, TimestampMixin):
    name = fields.CharField(max_length=128, unique=True, description="模型名称", index=True)
    desc = fields.CharField(max_length=1000, unique=True, description="模型描述", index=True)
    frame = fields.CharField(max_length=50, null=True, description="模型框架，如YOLOv8,paddle", index=True)
    task = fields.CharEnumField(IntelligenceModelTaskType,description="任务类型 是属于分类还是其他的", index=True)
    labels = fields.JSONField(max_length=1000, description="目标检测或分类标签 {id:0,name:'',}")
    is_quad = fields.BooleanField(default=False, description="如果task 是多边形或pose的模型有用，可以设置为True")
    class Meta:
        table = "intelligence_model"

class Datasets(BaseModel, TimestampMixin):
    """
    数据集
    """

    name = fields.CharField(max_length=128, unique=True, description="名称",index=True)
    desc = fields.CharField(max_length=1000, null=True, description="描述")
    file_type = fields.CharEnumField(enum_type=DataTypeEnum, default=DataTypeEnum.IMAGE, description="数据文件类型",index=True)
    files: fields.ReverseRelation["File"]
    tasks: fields.ReverseRelation["Job_Task"]  # 多对一

    # 计算关联文件数量数
    def file_count(self) -> int:
        return len(self.files)

    # 计算关联任务数量数
    def task_count(self) -> int:
        return len(self.tasks)

    class Meta:
        table = "datasets"
        table_description = "数据集"
        # indexes = ("name", "data_type")
        
class Job_Task(BaseModel, TimestampMixin):
    name = fields.CharField(max_length=128, unique=True, description="名称")
    desc = fields.CharField(max_length=1000, null=True, description="描述")
    status = fields.IntEnumField(enum_type=TaskStatusEnum, default=TaskStatusEnum.OPEN, description="是否属于打开关闭交付状态")
    datasets: fields.ForeignKeyNullableRelation["Datasets"] = fields.ForeignKeyField(
        "models.Datasets", related_name="tasks", null=True
    )
    model: fields.ForeignKeyNullableRelation["IntelligenceModel"] = fields.ForeignKeyField(
        "models.IntelligenceModel", related_name="tasks", null=True
    )  # 多对一
    auto_method_name = fields.CharField(max_length=1000, null=True, description="自动注释方法")  # 多对一 AI注释方法 通过这个名称调用方法，方法在ai_annotation.py中定义
    class Meta:
        table = "file"

class File(BaseModel, TimestampMixin):
    """
    文件信息
    """

    name = fields.CharField(max_length=128, unique=True, description="文件名")
    path = fields.CharField(max_length=255, unique=True, description="存储路径")
    datasets: fields.ForeignKeyNullableRelation["Datasets"] = fields.ForeignKeyField(
        "models.Datasets", related_name="files", null=True
    )

    class Meta:
        table = "file_infos"
        table_description = "文件信息"
        indexes = (
            "name",
            "path",
        )
        


class Annotation(BaseModel, TimestampMixin):
    status = fields.CharEnumField(
        enum_type=AnotationStatusEnum, default=AnotationStatusEnum.UNASSIGNED, description="任务状态"
    )
    annotation = fields.JSONField(default={}, description="标注数据")
    worker: fields.ForeignKeyNullableRelation["User"] = fields.ForeignKeyField(
        "models.User",
        on_delete=fields.OnDelete.CASCADE,
        related_name="task_data",
        default=None,
        null=True,
    )  # 一对一
    file: fields.ForeignKeyNullableRelation["File"] = fields.ForeignKeyField(
        "models.File",
        on_delete=fields.OnDelete.CASCADE,
        related_name="annotations",
    )  # 多对一
    task: fields.ForeignKeyNullableRelation["Job_Task"] = fields.ForeignKeyField(
        "models.Job_Task",
        on_delete=fields.OnDelete.CASCADE,
        related_name="annotations",
    )  # 多对一
    data_partition = fields.IntEnumField(enum_type=DataPartitionEnum, description="数据划分集(0=train,1=val,2=test)", default=None, null=True)
    comment: fields.OneToOneNullableRelation["AnotationComment"]
    class Meta:
        table = "annotation"


class AnotationComment(BaseModel, TimestampMixin):
    file_info: fields.OneToOneRelation["Annotation"] = fields.ForeignKeyField(
        "models.Annotation",
        default=None,
        on_delete=fields.OnDelete.CASCADE,
        related_name="comment",
    )  # 一对一
    annotation = fields.JSONField(default={}, description="评价")

# region 模型管理
# class Model_Pipeline(BaseModel, TimestampMixin):
#     name = fields.CharField(max_length=20, unique=True, description="角色名称", index=True)
#     desc = fields.CharField(max_length=500, null=True, description="角色描述")
#     models = fields.ManyToManyField("models.IntelligenceModel", related_name="相关模型")
#     pipeline = fields.JSONField(max_length=255, description="""
# 产线 [外框定位模型>分割>文本定位模型>分割>根据不同要素不同OCR模型],输出结果, 
# 节点: {
#     node_id:, 
#     model:, function: ,
#     input:, output:,
#     next:[{output:,next:node_id}, 指定那些输出内容到那个节点]
# }
# """)
#     class Meta:
#         table = "Model_Pipeline"
        
# class Model_Test_Result(BaseModel, TimestampMixin):
#     name = fields.CharField(max_length=20, unique=True, description="角色名称", index=True)
#     desc = fields.CharField(max_length=500, null=True, description="角色描述")
#     model = fields.ManyToManyField("models.IntelligenceModel", related_name="model_test_results")
#     test_files = fields.JSONField(max_length=255, description="测试注释文件 {id:0,name:'',}", index=True)
#     results = fields.JSONField(max_length=255, description="测试结果json {id:0,name:'',}", index=True)
#     class Meta:
#         table = "Model_Test_Result"
        
# class Pipeline_Test_Result(BaseModel, TimestampMixin):
#     name = fields.CharField(max_length=20, unique=True, description="角色名称", index=True)
#     desc = fields.CharField(max_length=500, null=True, description="角色描述")
#     model = fields.ManyToManyField("models.IntelligenceModel", related_name="model_test_results")
#     test_file = fields.JSONField(max_length=255, description="测试注释文件 {id:0,name:'',}", index=True)
#     results = fields.JSONField(max_length=255, description="测试结果json {id:0,name:'',}", index=True)
#     class Meta:
#         table = "Pipeline_Test_Result"

# class Model_Training(BaseModel, TimestampMixin):
#     name = fields.CharField(max_length=20, unique=True, description="数据模型", index=True)
#     desc = fields.CharField(max_length=500, null=True, description="角色描述")
#     model = fields.ManyToManyField("models.IntelligenceModel", related_name="model_test_results")
#     logs = fields.CharField(max_length=255, description="训练日志", index=True)
#     results = fields.JSONField(max_length=255, description="训练结果json {id:0,name:'',}", index=True)
#     class Meta:
#         table = "model_training"
# endregion

